// SPDX-License-Identifier: GPL-2.0

#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <linux/module.h>
#include <linux/printk.h>
#include <linux/vmalloc.h>
#include <crypto/hash.h>
#include <linux/highmem.h>
#include <linux/kallsyms.h>
#include <linux/rculist.h>
#include <linux/timer.h>
#include <linux/sched.h>
#include <linux/kthread.h>
#include <linux/jiffies.h>
#include <linux/utsname.h>
#include <linux/sched/signal.h>
#include <linux/limits.h>
#include <linux/sched/mm.h>
#include <linux/list.h>

#include "dim_core_measure.h"
#include "dim_core_measure_log.h"
#include "dim_core_policy.h"
#include "dim_core_dynamic_baseline.h"
#include "dim_core_static_baseline.h"
#include "dim_core_status.h"
#include "dim_core_sig.h"
#include "dim_core_mem.h"
#include "dim_baseline_node.h"

#include "../adapter/dim_adapter.h"
#include "../common/dim_hash.h"
#include "../common/dim_pages.h"
#include "../common/dim_calculate.h"
#include "../common/dim_utils.h"

#define DIM_RECURSION_TSK_MAX 64

/* get kernel module list */
extern struct list_head *get_kernel_modules(void);

/* variable for periodic measurement */
static struct task_struct *dim_measure_task;
static struct timer_list dim_measure_timer;
static unsigned long dim_measure_interval;
static unsigned int dim_measure_tampered_action;

static int dim_measure_user_tampered_action_kill(struct task_struct *tsk,
				unsigned int recurCount)
{
	int ret = 0;
	unsigned int count = recurCount;
	struct task_struct *task = NULL;

	if (tsk == NULL)
		return -EINVAL;

	if (count >= DIM_RECURSION_TSK_MAX) {
		pr_warn("number of subProc recursions oexceeds the maxValue %u\n",
				DIM_RECURSION_TSK_MAX);
		return 0;
	}

	if (tsk->tgid == 1)
		return 0;

	list_for_each_entry(task, &tsk->children, sibling) {
		count++;
		ret = dim_measure_user_tampered_action_kill(task, count);
		if (ret != 0)
			pr_warn("fail to tampered action recur pid %d\n", task->tgid);
		count--;
	}

	send_sig(SIGKILL, tsk, 1);

	return 0;
}

static void dim_measure_user_tampered_action(enum dim_core_policy_action action,
				struct task_struct *tsk)
{
	int ret;

	if (tsk == NULL || dim_measure_tampered_action == 0)
		return;

	switch (action) {
	case ACTION_KILL:
		read_lock(&tasklist_lock);
		ret = dim_measure_user_tampered_action_kill(tsk, 0);
		read_unlock(&tasklist_lock);
		if (ret < 0)
			pr_warn("fail to tampered kill proc\n");
		break;
	default:
		break;
	}

}

static int dim_measure_user_trampoline(struct mm_struct *mm,
				       struct vm_area_struct *vma,
				       char *path, struct task_struct *tsk)
{
	int ret, hash_len;
	struct dim_baseline_node *find_node = NULL;
	size_t name_len = strlen(path) + strlen(DIM_TRAMPOLINE_NAME_SUFFIX);
	char *name_buf = NULL;
	char digest_out[MAX_DIGEST_SIZE];

	if (name_len > 65536U)
		return -ENAMETOOLONG;

	find_node = dim_dynamic_baseline_search(path,
		DIM_BASELINE_TRAMPOLINE);
	if (find_node == NULL)
		return 0; /* no dynamic baseline, needn't measure */

	name_buf = kmalloc(name_len + 1, GFP_KERNEL);
	if (name_buf == NULL)
		return -ENOMEM;
	snprintf(name_buf, name_len + 1, "%s%s",
		path, DIM_TRAMPOLINE_NAME_SUFFIX);

	/* calculate and compare digest */
	hash_len = dim_hash_get_digest_len(find_node->algo);
	ret = dim_calc_user_task_trampoline_digest(mm, vma,
		find_node->algo, digest_out);
	if (ret != hash_len) {
		pr_warn("fail to calculate digest of trampoline of %s: %d\n",
			path, ret);
		ret = -EINVAL;
	} else if (dim_digest_compare(find_node->data, digest_out,
		hash_len) != 0) {
		dim_measure_baseline_record(name_buf, "tampered!",
			digest_out, find_node->algo, 0);
		dim_measure_user_tampered_action(find_node->action, tsk);
	}

	kfree(name_buf);
	return ret < 0 ? ret : 0;
}

static int dim_measure_user_text(struct mm_struct *mm,
				 struct vm_area_struct *vma,
				 char *path, struct task_struct *tsk)
{
	int ret, hash_len;
	struct dim_baseline_node *find_node = NULL;
	char digest_out[MAX_DIGEST_SIZE];

	/* get static baseline data and record */
	find_node = dim_dynamic_baseline_search(path, DIM_BASELINE_USER);
	if (find_node == NULL)
		return 0; /* no dynamic baseline, needn't measure */

	/* calculate and compare digest */
	hash_len = dim_hash_get_digest_len(find_node->algo);
	ret = dim_calc_user_task_text_digest(mm, vma,
		find_node->algo, digest_out);
	if (ret != hash_len) {
		pr_warn("fail to calculate digest of %s: %d\n", path, ret);
		return ret < 0 ? ret : -EINVAL;
	}

	if (dim_digest_compare(find_node->data, digest_out, hash_len) != 0) {
		dim_measure_baseline_record(path, "tampered!",
			digest_out, find_node->algo, 0);
		dim_measure_user_tampered_action(find_node->action, tsk);
	}
	return 0;
}

static int dim_measure_user_from_vma(struct mm_struct *mm,
				     struct vm_area_struct *vma,
				     char *path, struct task_struct *tsk)
{
	int ret;

	ret = dim_measure_user_trampoline(mm, vma, path, tsk);
	if (ret < 0) {
		pr_warn("fail to create trampo dynamic baseline of %s: %d\n",
			path, ret);
		return ret == -EINTR ? -EINTR : 0;
	}

	ret = dim_measure_user_text(mm, vma, path, tsk);
	if (ret < 0) {
		pr_warn("fail to create dynamic baseline of %s: %d\n",
			path, ret);
		return ret == -EINTR ? -EINTR : 0;
	}

	return 0;
}

struct measure_tsk {
	struct task_struct *t;
	struct list_head list;
};

static void add_to_measure_list(struct list_head *mtsks,
				struct task_struct *tsk)
{
	struct measure_tsk *mtsk = NULL;

	mtsk = kmalloc(sizeof(struct measure_tsk), GFP_ATOMIC);
	if (mtsk == NULL)
		return;

	mtsk->t = tsk;
	list_add_tail(&mtsk->list, mtsks);
}

static int dim_measure_user(void)
{
	struct task_struct *tsk = NULL;
	struct vm_area_struct *vma = NULL;
	struct mm_struct *mm = NULL;
	char *path = NULL;
	char *path_buf = NULL;
	struct file *vm_file_prev = NULL;
	struct list_head mtsks;
	struct measure_tsk *pos = NULL;
	struct measure_tsk *next = NULL;
	int ret;

	INIT_LIST_HEAD(&mtsks);

	path_buf = kmalloc(PATH_MAX, GFP_KERNEL);
	if (path_buf == NULL)
		return -ENOMEM;

	/* get all tasks, add every task's usage and add to measure task list */
	read_lock(&tasklist_lock);
	for_each_process(tsk) {
		get_task_struct(tsk);
		add_to_measure_list(&mtsks, tsk);
	}
	read_unlock(&tasklist_lock);

	/* do measure from measure task list */
	list_for_each_entry_safe(pos, next, &mtsks, list) {
		tsk = pos->t;
		mm = get_task_mm(tsk);
		if (mm == NULL) {
			put_task_struct(tsk);
			kfree(pos);
			continue;
		}

		down_read(&mm->mmap_lock);
		for (vma = mm->mmap; vma != NULL; vma = vma->vm_next) {
			if (!(vma->vm_flags & VM_READ) ||
			    !(vma->vm_flags & VM_EXEC) ||
			    (vma->vm_flags & VM_WRITE) ||
			    dim_get_vma_file(vma) == NULL ||
			    dim_get_vma_file(vma) == vm_file_prev)
				continue;

			vm_file_prev = dim_get_vma_file(vma);

			/* get elf file path from vma */
			path = d_path(&vm_file_prev->f_path,
				path_buf, PATH_MAX);
			if (IS_ERR(path)) {
				pr_warn("fail to get vm file: %ld\n",
					PTR_ERR(path));
				continue;
			}

			ret = dim_measure_user_from_vma(mm, vma, path, tsk);
			if (ret < 0)
				break;
		}
		up_read(&mm->mmap_lock);
		mmput(mm);
		put_task_struct(tsk);
		kfree(pos);
	}

	kfree(path_buf);
	return 0;
}

static int dim_measure_ko(void)
{
	struct module *mod = NULL;
	struct list_head *head = get_kernel_modules();
	char digest_out[MAX_DIGEST_SIZE];
	int ret, hash_len;
	char ko_name[NAME_MAX * 2 + 1];
	struct dim_baseline_node *find_node = NULL;

	if (head == NULL)
		return -EFAULT;

	mutex_lock(&module_mutex);
	list_for_each_entry_rcu(mod, head, list) {
		if (mod->state == MODULE_STATE_UNFORMED ||
		    mod->state != MODULE_STATE_LIVE)
			continue;

		if (strlen(init_uts_ns.name.release) +
			strlen(mod->name) + 1 >= sizeof(ko_name)) {
			pr_warn("overlength of %s or its kernel release!\n",
				ko_name);
			continue;
		}
		snprintf(ko_name, sizeof(ko_name), "%s/%s",
			init_uts_ns.name.release, mod->name);

		find_node = dim_dynamic_baseline_search(ko_name,
			DIM_BASELINE_KERNEL);
		if (find_node == NULL)
			continue;

		hash_len = dim_hash_get_digest_len(dim_hash_get_cur_index());
		ret = dim_calc_ko_digest(mod,
			dim_hash_get_cur_index(), digest_out);
		if (ret != hash_len) {
			pr_warn("fail to calculate %s digest: %d\n",
				ko_name, ret);
			continue;
		}

		if (dim_digest_compare(find_node->data,
			digest_out, hash_len) != 0) {
			pr_warn("ko dynamic baseline not match: %s\n",
				ko_name);
			dim_measure_baseline_record(ko_name, "tampered!",
				digest_out, dim_hash_get_cur_index(), 0);
			continue;
		}
	}
	mutex_unlock(&module_mutex);
	return 0;
}

static int dim_measure_kernel(void)
{
	struct dim_baseline_node *find_node = NULL;
	char digest_out[MAX_DIGEST_SIZE];
	int ret, hash_len;
	char *kernel_rel = init_uts_ns.name.release;

	find_node = dim_dynamic_baseline_search(kernel_rel,
		DIM_BASELINE_KERNEL);
	if (find_node == NULL) {
		dim_measure_baseline_record(kernel_rel, "no dynamic baseline!",
			NULL, dim_hash_get_cur_index(), 0);
		return 0;
	}

	hash_len = dim_hash_get_digest_len(dim_hash_get_cur_index());
	ret = dim_calc_kernel_digest(dim_hash_get_cur_index(), digest_out);
	if (ret != hash_len) {
		pr_warn("fail to calculate kernel digest: %d\n", ret);
		return 0;
	}

	if (dim_digest_compare(find_node->data, digest_out, hash_len) != 0) {
		pr_warn("kernel dynamic baseline not match\n");
		dim_measure_baseline_record(kernel_rel, "tampered!", digest_out,
			dim_hash_get_cur_index(), 0);
	}
	return 0;
}

int dim_measure_baseline_build(void)
{
	int ret, status;

	if (dim_lock_running_status() != 0)
		return -EBUSY;

	status = dim_get_running_status();
	if (status != DIM_NO_BASELINE &&
	    status != DIM_ERROR &&
	    status != DIM_PROTECTED) {
		dim_unlock_running_status();
		return -EBUSY;
	}

	dim_set_running_status(DIM_BASELINE_RUNNING);
	pr_info("do dim baseline init");

	dim_measure_log_prepare_for_baseline();

	ret = dim_sig_init();
	if (ret)
		goto out;

	ret = dim_policy_load();
	if (ret)
		goto out;

	ret = dim_static_baseline_build();
	if (ret)
		goto out;

	ret = dim_dynamic_baseline_build();
	if (ret)
		goto out;

	ret = 0;
out:
	dim_policy_destroy();
	dim_static_baseline_destroy();
	if (ret != 0)
		dim_dynamic_baseline_destroy();
	dim_set_running_status(ret == 0 ? DIM_PROTECTED : DIM_ERROR);
	dim_unlock_running_status();
	return ret;
}

int dim_measure_all(void)
{
	int ret, status;

	if (dim_lock_running_status() != 0)
		return -EBUSY;

	status = dim_get_running_status();
	if (status == DIM_NO_BASELINE || status == DIM_ERROR) {
		/* if no baseline, do baseline */
		dim_unlock_running_status();
		return dim_measure_baseline_build();
	}

	if (status != DIM_PROTECTED) {
		dim_unlock_running_status();
		return -EBUSY;
	}

	dim_set_running_status(DIM_MEASURE_RUNNING);
	pr_info("do dim measure");

	ret = dim_measure_user();
	if (ret)
		goto out;

	ret = dim_measure_ko();
	if (ret)
		goto out;

	ret = dim_measure_kernel();
	if (ret)
		goto out;

	ret = 0;
out:
	dim_set_running_status(ret == 0 ? DIM_PROTECTED : DIM_ERROR);
	dim_unlock_running_status();
	return ret;
}

void dim_measure_baseline_destroy(void)
{
	dim_dynamic_baseline_destroy();
	dim_static_baseline_destroy();
	dim_policy_destroy();
}

static void dim_measure_timer_handle(struct timer_list *timer)
{
	if (dim_measure_task == NULL)
		return;

	wake_up_process(dim_measure_task);
	if (dim_measure_interval > 0)
		mod_timer(&dim_measure_timer, jiffies +
			(HZ * 60 * dim_measure_interval));
}

static int dim_measure_task_handle(void *data)
{
	set_current_state(TASK_INTERRUPTIBLE);
	schedule();

	while (!kthread_should_stop()) {
		(void)dim_measure_all();
		set_current_state(TASK_INTERRUPTIBLE);
		schedule();
	}
	return 0;
}

int dim_measure_timer_init(void)
{
	int ret;

	dim_measure_task = kthread_run(dim_measure_task_handle,
		NULL, "dim_measure");
	if (IS_ERR(dim_measure_task)) {
		ret = PTR_ERR(dim_measure_task);
		dim_measure_task = NULL;
		pr_err("fail to create DIM periodic measurement task: %d\n",
			ret);
		return ret;
	}

	timer_setup(&dim_measure_timer, dim_measure_timer_handle, 0);
	return 0;
}

void dim_measure_timer_destroy(void)
{
	del_timer_sync(&dim_measure_timer);
	if (dim_measure_task != NULL)
		kthread_stop(dim_measure_task);
}

int dim_measure_timer_set_intervel(unsigned long m)
{
	if (m > DIM_INTERVAL_MAX || m < DIM_INTERVAL_MIN)
		return -EINVAL;

	pr_info("set measure interval: %lu", m);
	dim_measure_interval = m;
	if (m > 0)
		mod_timer(&dim_measure_timer, jiffies +
			(dim_measure_interval * 60 * HZ));

	return 0;
}

unsigned long dim_measure_timer_get_intervel(void)
{
	return dim_measure_interval;
}

int dim_measure_set_tampered_action(unsigned int m)
{
	pr_info("set tampered action: %u", m);
	dim_measure_tampered_action = m;

	return 0;
}

unsigned int dim_measure_get_tampered_action(void)
{
	return dim_measure_tampered_action;
}
